/*
 * Wazuh Vulnerability scanner - Scan Orchestrator
 * Copyright (C) 2015, Wazuh Inc.
 * Nov 23, 2023.
 *
 * This program is free software; you can redistribute it
 * and/or modify it under the terms of the GNU General Public
 * License (version 2) as published by the FSF - Free Software
 * Foundation.
 */

#ifndef _OS_DATA_CACHE_HPP
#define _OS_DATA_CACHE_HPP

#include "../policyManager/policyManager.hpp"
#include "cacheLRU.hpp"
#include "singleton.hpp"
#include "socketDBWrapper.hpp"
#include "wazuhDBQueryBuilder.hpp"
#include "wdbDataException.hpp"
#include <mutex>
#include <string>

/**
 * @brief Os structure.
 */
struct Os final
{
    std::string hostName;       ///< Hostname of the agent.
    std::string architecture;   ///< Architecture of the agent.
    std::string name;           ///< Name of operating system.
    std::string codeName;       ///< Code name of the operating system.
    std::string majorVersion;   ///< Major version of the operating system.
    std::string minorVersion;   ///< Minor version of the operating system.
    std::string patch;          ///< Patch of the operating system.
    std::string build;          ///< Build of the operating system.
    std::string platform;       ///< Platform of the operating system.
    std::string version;        ///< Version of the operating system.
    std::string release;        ///< Release of the operating system.
    std::string displayVersion; ///< Display version of the operating system.
    std::string sysName;        ///< System name of the operating system.
    std::string kernelVersion;  ///< Version of the kernel operating system.
    std::string kernelRelease;  ///< Release of the kernel operating system.
    std::string cpeName;        ///< CPE name of the operating system.
};

/**
 * @brief OsDataCache class.
 */
template<typename TSocketDBWrapper = SocketDBWrapper>
class OsDataCache final : public Singleton<OsDataCache<>>
{
private:
    LRUCache<std::string, Os> m_osData {PolicyManager::instance().getOsdataLRUSize()};
    std::mutex m_mutex;

    Os getOsDataFromWdb(const std::string& agentId) const
    {
        nlohmann::json response;
        try
        {
            TSocketDBWrapper::instance().query(WazuhDBQueryBuilder::builder().agentGetOsInfoCommand(agentId).build(),
                                               response);
        }
        catch (const SocketDbWrapperException& e)
        {
            throw WdbDataException(e.what(), agentId);
        }
        catch (const std::exception& e)
        {
            throw std::runtime_error("Unable to retrieve OS data from Wazuh-DB (agent " + agentId +
                                     "). Reason: " + e.what());
        }

        if (response.empty())
        {
            throw WdbDataException("Empty OS data from Wazuh-DB.", agentId);
        }

        Os osData;
        const auto& data = response.at(0);

        osData.hostName = data.value("hostname", "");
        osData.architecture = data.value("architecture", "");
        osData.name = data.value("os_name", "");
        osData.codeName = data.value("os_codename", "");
        osData.majorVersion = data.value("os_major", "");
        osData.minorVersion = data.value("os_minor", "");
        osData.patch = data.value("os_patch", "");
        osData.build = data.value("os_build", "");
        osData.platform = data.value("os_platform", "");
        osData.version = data.value("os_version", "");
        osData.release = data.value("os_release", "");
        osData.displayVersion = data.value("os_display_version", "");
        osData.sysName = data.value("sysname", "");
        osData.kernelVersion = data.value("version", "");
        osData.kernelRelease = data.value("release", "");
        // CPE must be built, it isn't present at the Wazuh-DB response.

        return osData;
    }

public:
    /**
     * @brief This method returns the os data.
     * @param agentId agent id.
     * @param osData OS data to fill.
     *
     * @return True if the OS data was found in the cache, false in case of Wazuh-DB query
     */
    bool getOsData(const std::string& agentId, Os& osData)
    {
        const auto getValue = [this](const std::string& agentIdValue) -> std::optional<Os>
        {
            std::scoped_lock lock(m_mutex);
            if (auto value = m_osData.getValue(agentIdValue); value)
            {
                return *value;
            }
            return std::nullopt;
        };

        if (const auto optOsData = getValue(agentId); optOsData != std::nullopt)
        {
            osData = *optOsData;
            return true;
        }

        // This may throw an exception that will be captured by the caller method.
        osData = getOsDataFromWdb(agentId);
        return false;
    }

    /**
     * @brief This method sets the os data.
     *
     * @param agentId agent id.
     * @param osData os data.
     */
    void setOsData(const std::string& agentId, const Os& osData)
    {
        std::scoped_lock lock(m_mutex);
        m_osData.insertKey(agentId, osData);
    }
};
#endif // _OS_DATA_CACHE_HPP
